import numpy as np
import pandas as pd
from sklearn import svm
from sklearn.model_selection import train_test_split


class SVM:
    """
    支持向量机
    """
    def __init__(self, dataset, args):
        input_dim = args.get('input_dim', 9)
        output_dim = args.get('output_dim', 1)

        df = pd.read_csv(dataset, delimiter=',', header=None)

        self.data_X = df.values[:, :input_dim]
        self.data_y = df.values[:, -output_dim]

        # 数据分割
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.data_X, self.data_y, random_state=666)

    def run(self):
        # fit the model, don't regularize for illustration purposes
        clf = svm.SVC(kernel='linear', C=1000)
        clf.fit(self.X_train, self.y_train)

        predict = clf.predict(self.X_test)
        all_data = np.hstack((self.X_test, self.y_test.reshape(-1, 1), predict.reshape(-1, 1)))

        cols = ["fat", "protein", "snf", "acidity", "lead", "mercury", "arsenic", "chromium", "aflatoxins", "truth", "predict"]
        raw_data = [{k: v for k, v in zip(cols, row)} for row in all_data.tolist()]

        return {
            'TrainAccuracy': clf.score(self.X_train, self.y_train),
            'TestAccuracy': clf.score(self.X_test, self.y_test),
            'Truth': list(map(int, list(self.y_test))),
            'Predict': list(map(int, list(predict))),
            'RawData': raw_data,
            'Description': '支持向量机（SVM）是一类按监督学习方式对数据进行二元分类的广义线性分类器，其决策边界是对学习样本求解的最大边距超平面。'
                           'SVM使用铰链损失函数（hinge loss）计算经验风险并在求解系统中加入了正则化项以优化结构风险，是一个具有稀疏性和稳健性的分类器。'
                           'SVM可以通过核方法（kernel method）进行非线性分类，是常见的核学习方法之一。'
        }
        # # 模型测试
        # with open('../../result/svm_%s.json' % self.dataset, 'w', encoding='utf8') as f:
        #     json.dump({'Train accuracy': clf.score(self.X_train, self.y_train),
        #                'Test accuracy': clf.score(self.X_test, self.y_test),
        #                'Truth': list(map(int, list(self.data_y))),
        #                'Predict': list(map(int, list(clf.predict(self.data_X))))},
        #               f)
        #
        # np.savetxt('../../result/svm_%s' % self.dataset, np.c_[self.data_X, clf.predict(self.data_X)], fmt='%.06f')


if __name__ == '__main__':
    print(SVM('../../data/sterilizedmilk_for_classfication.csv', {}).run())
